package com.seo.async;

import java.io.IOException;
import java.net.InetSocketAddress;
import java.net.SocketAddress;
import java.nio.channels.SelectionKey;
import java.nio.channels.Selector;
import java.nio.channels.ServerSocketChannel;
import java.nio.channels.SocketChannel;
import java.util.ArrayList;
import java.util.Collection;
import java.util.HashMap;
import java.util.Iterator;
import java.util.Map;

import org.apache.log4j.Logger;

public class AsyncSocket {

	private static final int DEFAULT_BACKLOG = 64;

	private Selector selector;

	private Map<SocketChannel, RegistrationInfo> sockets = new HashMap<SocketChannel, RegistrationInfo>();

	private Collection<TimerInfo> timers = new ArrayList<TimerInfo>();

	private int readTimeout = 10000;

	private int writeTimeout = 10000;

	private SocketChannel socketChannel;

	private static final Logger LOGGER = Logger.getLogger(AsyncSocket.class);

	private AsyncSocket(Selector selector) {
		this.selector = selector;
	}

	public int getReadTimeout() {
		return readTimeout;
	}

	public void setReadTimeout(int readTimeout) {
		this.readTimeout = readTimeout;
	}

	public int getWriteTimeout() {
		return writeTimeout;
	}

	public void setWriteTimeout(int writeTimeout) {
		this.writeTimeout = writeTimeout;
	}

	public static AsyncSocket create() throws IOException {
		return new AsyncSocket(Selector.open());
	}

	public void loop(final AsyncSocketListener listener) {
		listener.init(this);
		addTimer(1000, new Timer() {
			@Override
			public void onTimer() {
				checkTimeouts(listener);
			}
		});

		while (true) {
			try {
				internalLoop(listener);
			} catch (Exception e) {
				LOGGER.error(e.getMessage(), e);
				try {
					listener.onError(socketChannel, getContext(socketChannel));
				} catch (Exception e1) {
					LOGGER.error(e.getMessage(), e1);
				}
			} finally {
				invokeTimers();
			}
		}
	}

	private void invokeTimers() {
		for (TimerInfo timerInfo : timers) {
			if (timerInfo.nextInvoke <= System.currentTimeMillis()) {
				try {
					timerInfo.timer.onTimer();
				} catch (Exception e) {
					LOGGER.error(e.getMessage(), e);
				}
				timerInfo.nextInvoke = timerInfo.interval + System.currentTimeMillis();
			}
		}
	}

	private void checkTimeouts(AsyncSocketListener listener) {
		Iterator<SocketChannel> it = sockets.keySet().iterator();
		while (it.hasNext()) {
			SocketChannel socketChannel = it.next();
			RegistrationInfo registrationInfo = sockets.get(socketChannel);
			if (registrationInfo.isRegisteredRead() && (registrationInfo.getLastRead() + readTimeout <= System.currentTimeMillis())) {
				try {
					listener.onReadTimeout(socketChannel, getContext(socketChannel));
					return;
				} finally {
					close(socketChannel);
				}
			}
			if (registrationInfo.isRegisteredWrite() && (registrationInfo.getLastWrite() + writeTimeout <= System.currentTimeMillis())) {
				try {
					listener.onWriteTimeout(socketChannel, getContext(socketChannel));
					return;
				} finally {
					close(socketChannel);
				}
			}
		}
	}

	public void listen(SocketAddress address, int backlog) throws IOException {
		ServerSocketChannel channel = ServerSocketChannel.open();
		channel.socket().bind(address, backlog);
		channel.configureBlocking(false);
		channel.register(selector, SelectionKey.OP_ACCEPT);
	}

	public void listen(SocketAddress address) throws IOException {
		listen(address, DEFAULT_BACKLOG);
	}

	public SocketChannel connect(InetSocketAddress address, SocketContext socketContext) throws IOException {
		SocketChannel socketChannel = SocketChannel.open();
		socketChannel.configureBlocking(false);
		socketChannel.connect(address);
		sockets.put(socketChannel, new RegistrationInfo(socketContext));
		socketContext.setSocketChannel(socketChannel);
		return socketChannel;
	}

	public void close(SocketChannel socketChannel) {
		if (socketChannel != null) {
			sockets.remove(socketChannel);
			try {
				socketChannel.close();
			} catch (IOException e) {
				LOGGER.warn(e.getMessage(), e);
			}
		}
	}

	private void internalLoop(AsyncSocketListener listener) throws Exception {
		int count = selector.select(200);
		if (count < 0) {
			return;
		}

		Iterator<SelectionKey> it = selector.selectedKeys().iterator();
		while (it.hasNext()) {
			SelectionKey selectionKey = it.next();
			it.remove();
			if (!selectionKey.isValid()) {
				continue;
			}

			if (selectionKey.isAcceptable()) {
				ServerSocketChannel serverSocketChannel = (ServerSocketChannel) selectionKey.channel();
				socketChannel = serverSocketChannel.accept();
				socketChannel.configureBlocking(false);
				if (socketChannel != null) {
					RegistrationInfo registrationInfo = new RegistrationInfo();
					sockets.put(socketChannel, registrationInfo);
					Object context = listener.onAccept(socketChannel);
					registrationInfo.setContext(context);
				}
				continue;
			}

			if (selectionKey.isConnectable()) {
				SocketChannel socketChannel = (SocketChannel) selectionKey.channel();
				if (!socketChannel.finishConnect()) {
					selectionKey.cancel();
					listener.onError(socketChannel, getContext(socketChannel));
				}
			}

			if (selectionKey.isReadable()) {
				socketChannel = (SocketChannel) selectionKey.channel();
				RegistrationInfo registrationInfo = sockets.get(socketChannel);
				try {
					listener.onRead(socketChannel, registrationInfo.getContext());
				} finally {
					registrationInfo.setLastRead(System.currentTimeMillis());
				}
			}

			if (selectionKey.isValid() && selectionKey.isWritable()) {
				socketChannel = (SocketChannel) selectionKey.channel();
				RegistrationInfo registrationInfo = sockets.get(socketChannel);
				try {
					listener.onWrite(socketChannel, getContext(socketChannel));
				} finally {
					registrationInfo.setLastWrite(System.currentTimeMillis());
				}
			}
		}
	}

	private Object getContext(SocketChannel socketChannel) {
		RegistrationInfo registrationInfo = sockets.get(socketChannel);
		if (registrationInfo != null) {
			return registrationInfo.getContext();
		}
		return null;
	}

	public void subscribe(SocketChannel socketChannel, int options) throws IOException {
		RegistrationInfo registrationInfo = sockets.get(socketChannel);
		if (registrationInfo.getOptions() != options) {
			if (!registrationInfo.isRegisteredRead() && ((options & SelectionKey.OP_READ) != 0)) {
				registrationInfo.setLastRead(System.currentTimeMillis());
			}
			if (!registrationInfo.isRegisteredWrite() && ((options & SelectionKey.OP_WRITE) != 0)) {
				registrationInfo.setLastWrite(System.currentTimeMillis());
			}
			socketChannel.register(selector, options);
			registrationInfo.setOptions(options);
		}
	}

	public void unsubscribe(SocketChannel socketChannel) throws IOException {
		subscribe(socketChannel, 0);
	}

	public void addTimer(int interval, Timer timer) {
		timers.add(new TimerInfo(interval, timer));
	}
}
